package telex;
/*
 * QueryRouter.java
 *
 * Copyright (c) 2007-2008 Colin Bayer, Douglas Downey, Oren Etzioni,
 *   University of Washington Computer Science and Engineering
 * 
 * See the file "COPYING" in the root directory of this distribution for
 * full licensing information.
 */

import java.io.File;
import java.util.*;

public class QueryRouter {
	
	private static void flushQueries(HashMap<String, ArrayList<Query>> locToQuery,
			String aOutputDir) throws Exception {
		Iterator<String> locIt = locToQuery.keySet().iterator();
		while(locIt.hasNext()) {
			String loc = locIt.next();
                        String path = aOutputDir + File.separator + loc;
			IQFWriter w = new IQFWriter(path, false, true);
			Iterator<Query> qIt = locToQuery.get(loc).iterator();
			while(qIt.hasNext()) w.writeQuery(qIt.next());
			w.finish();
		}
	}
	
	//Returns the number of unique queries (*not* query doc pairs) processed
	//HACK: returns -1 if end of file is reached
	public static int route(String aInputPath, String aOutputDir, int maxQueriesInMemory,
			int skipQueries) {
		IQFReader ir;
		//HashSet<String> locs_all_files = new HashSet<String>();
		HashMap<String, ArrayList<Query>> locToQuery = new HashMap<String, ArrayList<Query>>();
		
		
		try {
			ir = new IQFReader(aInputPath);
		} catch (Exception e) {
			System.err.println("Error opening " + aInputPath + " for reading: " + e);
			return -1;
		}
		
		// scan the input file, recording queries and target locations
		// when memory limit is reached, output the queries to files
		int totalQueries = 0;
		int queryDocPairs = 0;
		try {
			Query q;
			while ((q = ir.nextQuery()) != null) {
				if(skipQueries > 0) {
					skipQueries--; continue;
				}
				for (String loc : q.getUserFeature("loc").split(",")) {
					if(!locToQuery.containsKey(loc))
						locToQuery.put(loc, new ArrayList<Query>());
					locToQuery.get(loc).add(q);
					queryDocPairs++;
				}
				totalQueries++;
				//for routed output, null out the location variable:
				q.setUserFeature("loc", "");
				if(++queryDocPairs >= maxQueriesInMemory) {
					break;
				}
			}
			flushQueries(locToQuery, aOutputDir);
			ir.close();
		} catch (Exception e) {
			System.err.println("Error processing queries from " + aInputPath + ": " + e);
			return -1;
		}
		ir.close();
		if(queryDocPairs < maxQueriesInMemory)
			return -1;
		return totalQueries;
	}
	
	static void banner() {
		System.out.print(
			"teLEX Query Router v. 0.1\n" +
			"Colin Bayer and Douglas Downey <{vogon,ddowney}@cs.washington.edu>\n" +
			"Based on an algorithm in 'Locating Named Entities in Web Text'\n" +
				"\t(Downey, Broadhead, Etzioni 2007, published in IJCAI 2007 Proceedings)\n" +
			"See README for a description of the Query Sorter and teLEX in general.\n" +
			"=====================================================================================\n"
		);
	}
	
	static void usage() {
		banner();
		System.err.print(
			"Usage: java QueryRouter [-fh n] input [output]\n" +
			"-fh n: limit number of query/doc records to hold in memory at once to n (default: 10,000,000)\n" +
			"\tinput: path to unsorted IQF, or directory containing one or more\n" +
				"\t\tIQF files to sort\n" +
			"\toutput: path to directory to output IQF file(s) in (default: input + \"-rt\")\n"
		);
		System.exit(255);
	}
	
	public static void main(String[] aArgs) {
		if (aArgs.length < 1) {
			usage();
		}
		
		String input_path = null, output_path = null;
		int handle_limit = 10000000;
		
		for (int i = 0; i < aArgs.length; i++) {
			boolean silence_pathname_warning = false;
			
			if (aArgs[i].equals("-fh")) {
				// -fh (file handle limit) option.
				try {
					handle_limit = Integer.parseInt(aArgs[i + 1]);
				} catch (Exception e) {
					System.err.println("Option '-fh' requires an integer argument.");
					usage();
				}
				
				i++;
			} else if (aArgs[i].charAt(0) == '-') {
				System.err.println("Unrecognized option '" + aArgs[i] + "'.");
				usage();
			} else {
				// pathname in arguments list.  assign it to the first of input_path or output_path,
				// or warn if both are already assigned.
				if (input_path == null) input_path = aArgs[i];
				else if (output_path == null) output_path = aArgs[i];
				else if (!silence_pathname_warning) {
					System.err.println("Warning: unused pathname argument ignored.");
					silence_pathname_warning = true;
				}
			}
		}
		
		ArrayList<String> files_to_process = new ArrayList<String>();
		
		banner();
		
		// check that at least one of the two pathnames was supplied.
		if (input_path == null) {
			System.out.println("No input path specified.");
			usage();
		} else {
			File f = new File(input_path);
			boolean input_is_dir = false;
			
			// detect directory-ness of doc_path.	
			if (!f.exists()) {
				System.err.println("Input file does not exist.");
				System.exit(1);
			} else if (f.isDirectory()) {
				System.out.println("<== Input directory: " + f);
				input_is_dir = true;
			} else if (f.isFile()) {
				System.out.println("<== Input file: " + f);
				input_is_dir = false;
			} else {
				System.err.println("Input file exists, but is not a normal file or a directory.");
				System.exit(1);
			}
			
			// generate output path, if none provided.
			if (output_path == null) {
				if (input_is_dir) output_path = input_path + "-rt";
				else output_path = input_path + ".routed";
			}
			
			// check that the output path either doesn't exist, or is of the right type to overwrite.
			// create the output directory if we're writing a directory.
			f = new File(output_path);

			System.out.println("==> Output directory " + (f.exists() ? ": " : " (creating): ") + output_path);
			
			if (!f.exists()) {
				try {
					f.mkdirs();
				} catch (Exception e) {
					System.err.println("Error trying to create output directory: " + e);
					System.exit(1);
				}
			} else if (f.isFile()) {
				System.err.println("Output file already exists, but isn't a directory; refusing to overwrite it.");
				System.exit(1);
			}
			
			// generate list of input and output filenames to process.
			if (input_is_dir) {
				f = new File(input_path);
				
				String[] files_in_dir = f.list();
				
				if (files_in_dir == null) {
					System.err.println("Error getting list of files in input directory.  Halting.");
					System.exit(1);
				}
				
				for (String fname : files_in_dir) {
					files_to_process.add(input_path + "/" + fname);
				}
			} else {
				files_to_process.add(input_path);
			}
		}
		
		for (String in_file : files_to_process) {
			System.out.println("Routing queries from " + in_file + "...");
			int totalQueriesProcessed = 0;
			int lastQueriesProcessed = 0;
			
			while((lastQueriesProcessed = route(in_file, output_path, handle_limit, totalQueriesProcessed)) > 0) {
				totalQueriesProcessed += lastQueriesProcessed;
				System.out.println("Processed " + totalQueriesProcessed + " queries so far.");
			}
			System.out.println("Done.");
		}
		System.out.println("Routing completed successfully.");
	}
}

